/*
 * @FilePath: \DWARF2_V2\src\track\yolov8.h
 * @Description:
 * @Author: baike && baike93@163.com
 * @Date: 2024-03-09
 * @LastEditors: baike && baike93@163.com
 * @LastEditTime: 2024-03-23
 * Copyright (c) 2024 by Tinyphoton (Wuhan) Technology Co., LTD, All Rights
 * Reserved.
 */

// Cpp native
#include <fstream>
#include <random>
#include <string>
#include <vector>

// OpenCV / DNN / YoloV8
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/opencv.hpp>

#include "./rknn_demo_utils.h"
#include "./yolo.h"
#include "im2d.h"
#include "rga.h"
#include "rkmedia_buffer.h"
#include "rknn_api.h"


class YoloV8 {
 public:
  YoloV8(const std::string &rknnModelPath);
  int predict(const cv::Mat &input, detect_result_group_t &detect_result_group);
  int predict_with_nv12(MEDIA_BUFFER fram_mb, cv::Size frame_size,
                        detect_result_group_t &detect_result_group);

 private:
  cv::Mat preprocess(const cv::Mat &image, float &scale_w, float &scale_h);
  void preprocess_with_rga(const cv::Mat &image, void *dst_buffer,
                           float &scale_w, float &scale_h);
  void preprocess_with_rga(MEDIA_BUFFER fram_mb, cv::Size frame_size,
                           void *dst_buffer, float &scale_w, float &scale_h);
  int postprocess(int input_w, int input_h, float scale_w, float scale_h,
                  detect_result_group_t &detect_result_group);

  rknn_context yolov8_ctx_;
  int model_input_size_ = 640;
  MODEL_INFO model_info_;
  YOLO_INFO y_info_;
};